package dns

import (
	"context"
	"fmt"
	"strings"
	"time"

	"github.com/catroll/fastdns/internal/logger"
)

// LookupRecords 查询指定类型的 DNS 记录
func (r *Resolver) LookupRecords(ctx context.Context, domain string, recordType RecordType) ([]any, error) {
	logger.Info("开始 DNS 查询", "domain", domain, "type", recordType)
	startTime := time.Now()

	// 检查缓存
	if r.cache != nil {
		if records, found := r.cache.GetRecords(domain, recordType); found {
			logger.Info("DNS 记录缓存命中", "domain", domain, "type", recordType, "records", len(records), "duration", time.Since(startTime))
			return records, nil
		}
		logger.Debug("DNS 记录缓存未命中", "domain", domain, "type", recordType)
	}

	var records []any

	switch recordType {
	case RTypeA, RTypeAAAA:
		// 查询 A 或 AAAA 记录（IP 地址）
		logger.Debug("执行 IP 地址查询", "domain", domain, "type", recordType)
		ips, err := r.LookupIP(ctx, domain)

		// 如果仍然失败，返回错误
		if err != nil {
			logger.Error("DNS 查询失败", "domain", domain, "type", recordType, "error", err, "duration", time.Since(startTime))
			return nil, err
		}

		// 如果没有记录（可能是智能模式下添加了 CNAME 记录），则添加 IP 记录
		if len(records) == 0 {
			for _, ip := range ips {
				rType := RTypeA
				if ip.To4() == nil {
					rType = RTypeAAAA
				}

				records = append(records, DNSRecord{
					Type:  rType,
					Name:  domain,
					Value: ip.String(),
					TTL:   int(r.config.DNS.Cache.TTL.Seconds()),
				})
			}
		}

	case RTypeCNAME:
		// 查询 CNAME 记录
		logger.Debug("执行 CNAME 查询", "domain", domain)
		cname, err := r.lookupCNAME(ctx, domain)
		if err != nil {
			logger.Error("DNS 查询失败", "domain", domain, "type", recordType, "error", err, "duration", time.Since(startTime))
			return nil, err
		}

		record := DNSRecord{
			Type:  RTypeCNAME,
			Name:  domain,
			Value: cname,
			TTL:   int(r.config.DNS.Cache.TTL.Seconds()),
		}
		records = []any{record}

	case RTypeMX:
		// 查询 MX 记录
		logger.Debug("执行 MX 查询", "domain", domain)
		mxRecords, err := r.lookupMX(ctx, domain)
		if err != nil {
			logger.Error("DNS 查询失败", "domain", domain, "type", recordType, "error", err, "duration", time.Since(startTime))
			return nil, err
		}

		for _, mx := range mxRecords {
			// 将优先级信息添加到值中，因为 DNSRecord 没有 Priority 字段
			value := fmt.Sprintf("%d %s", mx.Priority, mx.Host)
			records = append(records, DNSRecord{
				Type:  RTypeMX,
				Name:  domain,
				Value: value,
				TTL:   int(r.config.DNS.Cache.TTL.Seconds()),
			})
		}

	case RTypeTXT:
		// 查询 TXT 记录
		logger.Debug("执行 TXT 查询", "domain", domain)
		txtRecords, err := r.lookupTXT(ctx, domain)
		if err != nil {
			logger.Error("DNS 查询失败", "domain", domain, "type", recordType, "error", err, "duration", time.Since(startTime))
			return nil, err
		}

		for _, txt := range txtRecords {
			records = append(records, DNSRecord{
				Type:  RTypeTXT,
				Name:  domain,
				Value: txt,
				TTL:   int(r.config.DNS.Cache.TTL.Seconds()),
			})
		}

	case RTypeNS:
		// 查询 NS 记录
		logger.Debug("执行 NS 查询", "domain", domain)
		nsRecords, err := r.lookupNS(ctx, domain)
		if err != nil {
			logger.Error("DNS 查询失败", "domain", domain, "type", recordType, "error", err, "duration", time.Since(startTime))
			return nil, err
		}

		for _, ns := range nsRecords {
			records = append(records, DNSRecord{
				Type:  RTypeNS,
				Name:  domain,
				Value: ns,
				TTL:   int(r.config.DNS.Cache.TTL.Seconds()),
			})
		}

	default:
		err := fmt.Errorf("不支持的记录类型: %s", recordType)
		logger.Error("DNS 查询失败", "domain", domain, "type", recordType, "error", err, "duration", time.Since(startTime))
		return nil, err
	}

	// 更新缓存
	if r.cache != nil {
		r.cache.SetRecords(domain, recordType, records)
		logger.Debug("DNS 记录缓存更新", "domain", domain, "recordType", recordType, "records", records)
	}

	logger.Info("DNS 查询完成", "domain", domain, "recordType", recordType, "records", records, "duration", time.Since(startTime))
	return records, nil
}

// lookupCNAME 查询 CNAME 记录
func (r *Resolver) lookupCNAME(ctx context.Context, domain string) (string, error) {
	// 使用第一个健康的上游服务器查询
	r.upstreamsMu.RLock()
	defer r.upstreamsMu.RUnlock()

	for _, upstream := range r.upstreams {
		if !upstream.healthy {
			continue
		}

		cname, err := upstream.client.LookupCNAME(ctx, domain)
		if err != nil {
			logger.Debug("CNAME 查询失败",
				"domain", domain,
				"server", upstream.Address,
				"error", err)
			continue
		}

		// 移除末尾的点
		cname = strings.TrimSuffix(cname, ".")
		return cname, nil
	}

	return "", fmt.Errorf("所有上游 DNS 服务器查询 CNAME 失败")
}

// lookupMX 查询 MX 记录
func (r *Resolver) lookupMX(ctx context.Context, domain string) ([]MXRecord, error) {
	// 使用第一个健康的上游服务器查询
	r.upstreamsMu.RLock()
	defer r.upstreamsMu.RUnlock()

	for _, upstream := range r.upstreams {
		if !upstream.healthy {
			continue
		}

		mxs, err := upstream.client.LookupMX(ctx, domain)
		if err != nil {
			logger.Debug("MX 查询失败",
				"domain", domain,
				"server", upstream.Address,
				"error", err)
			continue
		}

		// 转换为我们的 MXRecord 格式
		records := make([]MXRecord, 0, len(mxs))
		for _, mx := range mxs {
			// 移除末尾的点
			host := strings.TrimSuffix(mx.Host, ".")
			records = append(records, MXRecord{
				Host:     host,
				Priority: int(mx.Pref),
				TTL:      int(r.config.DNS.Cache.TTL.Seconds()),
			})
		}

		return records, nil
	}

	return nil, fmt.Errorf("所有上游 DNS 服务器查询 MX 失败")
}

// lookupTXT 查询 TXT 记录
func (r *Resolver) lookupTXT(ctx context.Context, domain string) ([]string, error) {
	// 使用第一个健康的上游服务器查询
	r.upstreamsMu.RLock()
	defer r.upstreamsMu.RUnlock()

	for _, upstream := range r.upstreams {
		if !upstream.healthy {
			continue
		}

		txts, err := upstream.client.LookupTXT(ctx, domain)
		if err != nil {
			logger.Debug("TXT 查询失败",
				"domain", domain,
				"server", upstream.Address,
				"error", err)
			continue
		}

		return txts, nil
	}

	return nil, fmt.Errorf("所有上游 DNS 服务器查询 TXT 失败")
}

// LookupRecordsFromRandomUpstream 从随机选择的 DNS 服务器查询记录
func (r *Resolver) LookupRecordsFromRandomUpstream(ctx context.Context, domain string, recordType RecordType) ([]any, error) {
	// 检查缓存
	if r.cache != nil {
		if records, found := r.cache.GetRecords(domain, recordType); found {
			logger.Info("DNS 记录缓存命中", "domain", domain, "type", recordType, "records", len(records))
			return records, nil
		}
		logger.Debug("DNS 记录缓存未命中", "domain", domain, "type", recordType)
	}

	// 锁定上游服务器列表
	r.upstreamsMu.RLock()
	defer r.upstreamsMu.RUnlock()

	// 筛选健康的上游服务器
	healthyUpstreams := make([]*UpstreamServer, 0)
	for _, upstream := range r.upstreams {
		if upstream.healthy {
			healthyUpstreams = append(healthyUpstreams, upstream)
		}
	}

	if len(healthyUpstreams) == 0 {
		return nil, fmt.Errorf("没有可用的健康 DNS 服务器")
	}

	// 随机选择一个健康的上游服务器
	randIndex := time.Now().UnixNano() % int64(len(healthyUpstreams))
	selectedUpstream := healthyUpstreams[randIndex]

	logger.Debug("随机选择 DNS 服务器", "address", selectedUpstream.Address, "domain", domain, "type", recordType)

	var records []any

	switch recordType {
	case RTypeA, RTypeAAAA:
		// 创建带超时的上下文
		timeoutCtx, cancel := context.WithTimeout(ctx, selectedUpstream.Timeout)
		defer cancel()

		// 查询 IP 地址
		start := time.Now()
		ips, err := selectedUpstream.client.LookupIP(timeoutCtx, "ip", domain)
		latency := time.Since(start)
		selectedUpstream.latency = latency

		if err != nil {
			logger.Debug("IP 查询失败", "domain", domain, "server", selectedUpstream.Address, "error", err)
			return nil, err
		}

		// 构建记录
		for _, ip := range ips {
			rType := RTypeA
			if ip.To4() == nil {
				rType = RTypeAAAA
			}

			records = append(records, DNSRecord{
				Type:  rType,
				Name:  domain,
				Value: ip.String(),
				TTL:   int(r.config.DNS.Cache.TTL.Seconds()),
			})
		}

	case RTypeCNAME:
		// 创建带超时的上下文
		timeoutCtx, cancel := context.WithTimeout(ctx, selectedUpstream.Timeout)
		defer cancel()

		// 查询 CNAME
		start := time.Now()
		cname, err := selectedUpstream.client.LookupCNAME(timeoutCtx, domain)
		latency := time.Since(start)
		selectedUpstream.latency = latency

		if err != nil {
			logger.Debug("CNAME 查询失败", "domain", domain, "server", selectedUpstream.Address, "error", err)
			return nil, err
		}

		// 移除末尾的点
		cname = strings.TrimSuffix(cname, ".")
		records = []any{DNSRecord{
			Type:  RTypeCNAME,
			Name:  domain,
			Value: cname,
			TTL:   int(r.config.DNS.Cache.TTL.Seconds()),
		}}

	case RTypeMX:
		// 创建带超时的上下文
		timeoutCtx, cancel := context.WithTimeout(ctx, selectedUpstream.Timeout)
		defer cancel()

		// 查询 MX
		start := time.Now()
		mxs, err := selectedUpstream.client.LookupMX(timeoutCtx, domain)
		latency := time.Since(start)
		selectedUpstream.latency = latency

		if err != nil {
			logger.Debug("MX 查询失败", "domain", domain, "server", selectedUpstream.Address, "error", err)
			return nil, err
		}

		// 构建记录
		for _, mx := range mxs {
			// 移除末尾的点
			host := strings.TrimSuffix(mx.Host, ".")
			// 将优先级信息添加到值中
			value := fmt.Sprintf("%d %s", mx.Pref, host)
			records = append(records, DNSRecord{
				Type:  RTypeMX,
				Name:  domain,
				Value: value,
				TTL:   int(r.config.DNS.Cache.TTL.Seconds()),
			})
		}

	case RTypeTXT:
		// 创建带超时的上下文
		timeoutCtx, cancel := context.WithTimeout(ctx, selectedUpstream.Timeout)
		defer cancel()

		// 查询 TXT
		start := time.Now()
		txts, err := selectedUpstream.client.LookupTXT(timeoutCtx, domain)
		latency := time.Since(start)
		selectedUpstream.latency = latency

		if err != nil {
			logger.Debug("TXT 查询失败", "domain", domain, "server", selectedUpstream.Address, "error", err)
			return nil, err
		}

		// 构建记录
		for _, txt := range txts {
			records = append(records, DNSRecord{
				Type:  RTypeTXT,
				Name:  domain,
				Value: txt,
				TTL:   int(r.config.DNS.Cache.TTL.Seconds()),
			})
		}

	case RTypeNS:
		// 创建带超时的上下文
		timeoutCtx, cancel := context.WithTimeout(ctx, selectedUpstream.Timeout)
		defer cancel()

		// 查询 NS
		start := time.Now()
		nss, err := selectedUpstream.client.LookupNS(timeoutCtx, domain)
		latency := time.Since(start)
		selectedUpstream.latency = latency

		if err != nil {
			logger.Debug("NS 查询失败", "domain", domain, "server", selectedUpstream.Address, "error", err)
			return nil, err
		}

		// 构建记录
		for _, ns := range nss {
			// 移除末尾的点
			host := strings.TrimSuffix(ns.Host, ".")
			records = append(records, DNSRecord{
				Type:  RTypeNS,
				Name:  domain,
				Value: host,
				TTL:   int(r.config.DNS.Cache.TTL.Seconds()),
			})
		}

	default:
		return nil, fmt.Errorf("不支持的记录类型: %s", recordType)
	}

	// 更新缓存
	if r.cache != nil && len(records) > 0 {
		r.cache.SetRecords(domain, recordType, records)
		logger.Debug("DNS 记录缓存更新", "domain", domain, "recordType", recordType, "records", records)
	}

	return records, nil
}

// lookupNS 查询 NS 记录
func (r *Resolver) lookupNS(ctx context.Context, domain string) ([]string, error) {
	// 使用第一个健康的上游服务器查询
	r.upstreamsMu.RLock()
	defer r.upstreamsMu.RUnlock()

	for _, upstream := range r.upstreams {
		if !upstream.healthy {
			continue
		}

		nss, err := upstream.client.LookupNS(ctx, domain)
		if err != nil {
			logger.Debug("NS 查询失败",
				"domain", domain,
				"server", upstream.Address,
				"error", err)
			continue
		}

		// 转换为字符串列表
		nsStrings := make([]string, 0, len(nss))
		for _, ns := range nss {
			// 移除末尾的点
			host := strings.TrimSuffix(ns.Host, ".")
			nsStrings = append(nsStrings, host)
		}

		return nsStrings, nil
	}

	return nil, fmt.Errorf("所有上游 DNS 服务器查询 NS 失败")
}
